{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b7973eef-7e5e-4ca4-844a-1d016f5a638b",
   "metadata": {},
   "source": [
    "# Building an Agent around a Query Pipeline\n",
    "\n",
    "In this cookbook we show you how to build an agent around a query pipeline.\n",
    "\n",
    "Agents offer the ability to do complex, sequential reasoning on top of any query DAG that you have setup. Conceptually this is also one of the ways you can add a \"loop\" to the graph.\n",
    "\n",
    "We show you two examples of agents you can implement:\n",
    "- a full ReAct agent that can do tool picking\n",
    "- a \"simple\" agent that adds a retry layer around a text-to-sql query engine."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "122dce22-6d3a-4d4a-a265-a6a3d3f90d26",
   "metadata": {},
   "source": [
    "## Setup Data\n",
    "\n",
    "We use the chinook database as sample data. [Source](https://www.sqlitetutorial.net/sqlite-sample-database/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "926b79ba-1868-46ca-bb7e-2bd3c907773c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
      "                                 Dload  Upload   Total   Spent    Left  Speed\n",
      "100  298k  100  298k    0     0  2327k      0 --:--:-- --:--:-- --:--:-- 2387k\n",
      "curl: (6) Could not resolve host: .\n",
      "Archive:  ./chinook.zip\n",
      "  inflating: chinook.db              \n"
     ]
    }
   ],
   "source": [
    "!curl \"https://www.sqlitetutorial.net/wp-content/uploads/2018/03/chinook.zip\" -O ./chinook.zip\n",
    "!unzip ./chinook.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "763a442a-cfcd-4e63-9121-e3a45dc3acff",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import SQLDatabase\n",
    "from sqlalchemy import (\n",
    "    create_engine,\n",
    "    MetaData,\n",
    "    Table,\n",
    "    Column,\n",
    "    String,\n",
    "    Integer,\n",
    "    select,\n",
    "    column,\n",
    ")\n",
    "\n",
    "engine = create_engine(\"sqlite:///chinook.db\")\n",
    "sql_database = SQLDatabase(engine)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0aacdcad-f0c1-40f4-b319-6c4cf3b309c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.query_pipeline import QueryPipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "419f97cf-63c1-456b-babd-b07d9ce4b937",
   "metadata": {},
   "source": [
    "### Setup Observability\n",
    "\n",
    "We setup Arize Phoenix for observability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "754f11d3-f053-46f7-acb4-ae8ee7d3fe07",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🌍 To view the Phoenix app in your browser, visit http://127.0.0.1:6006/\n",
      "📺 To view the Phoenix app in a notebook, run `px.active_session().view()`\n",
      "📖 For more information on how to use Phoenix, check out https://docs.arize.com/phoenix\n"
     ]
    }
   ],
   "source": [
    "# setup Arize Phoenix for logging/observability\n",
    "import phoenix as px\n",
    "import llama_index\n",
    "\n",
    "px.launch_app()\n",
    "llama_index.set_global_handler(\"arize_phoenix\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbbc38d8-2a78-4457-888d-cfe8960a9c6b",
   "metadata": {},
   "source": [
    "## Setup Text-to-SQL Query Engine / Tool\n",
    "\n",
    "Now we setup a simple text-to-SQL tool: given a query, translate text to SQL, execute against database, and get back a result. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d72eb79e-05b4-4260-b086-e837b01cce77",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.query_engine import NLSQLTableQueryEngine\n",
    "from llama_index.tools.query_engine import QueryEngineTool\n",
    "\n",
    "sql_query_engine = NLSQLTableQueryEngine(\n",
    "    sql_database=sql_database,\n",
    "    tables=[\"albums\", \"tracks\", \"artists\"],\n",
    "    verbose=True,\n",
    ")\n",
    "sql_tool = QueryEngineTool.from_defaults(\n",
    "    query_engine=sql_query_engine,\n",
    "    name=\"sql_tool\",\n",
    "    description=(\n",
    "        \"Useful for translating a natural language query into a SQL query\"\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af9d8d26-65c8-4788-a0ad-8d2448b23254",
   "metadata": {},
   "source": [
    "## Setup ReAct Agent Pipeline\n",
    "\n",
    "We now setup a ReAct pipeline for a single step using our Query Pipeline syntax. This is a multi-part process that does the following:\n",
    "1. Takes in agent inputs\n",
    "2. Calls ReAct prompt using LLM to generate next action/tool (or returns a response).\n",
    "3. If tool/action is selected, call tool pipeline to execute tool + collect response.\n",
    "4. If response is generated, get response.\n",
    "\n",
    "Throughout this we'll use a variety of agent-specific query components. Unlike normal query pipelines, these are specifically designed for query pipelines that are used in a `QueryPipelineAgentWorker`:\n",
    "- An `AgentInputComponent` that allows you to convert the agent inputs (Task, state dictionary) into a set of inputs for the query pipeline.\n",
    "- An `AgentFnComponent`: a general processor that allows you to take in the current Task, state, as well as any arbitrary inputs, and returns an output. In this cookbook we define a function component to format the ReAct prompt. However, you can put this anywhere.\n",
    "- An `CustomAgentComponent`: similar to `AgentFnComponent`, you can implement `_run_component` to define your own logic, with access to Task and state. It is more verbose but more flexible than `AgentFnComponent` (e.g. you can define init variables, and callbacks are in the base class).\n",
    "\n",
    "Note that any function passed into `AgentFnComponent` and `AgentInputComponent` MUST include `task` and `state` as input variables, as these are inputs passed from the agent. \n",
    "\n",
    "Note that the output of an agentic query pipeline MUST be `Tuple[AgentChatResponse, bool]`. You'll see this below."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a024201f-a39b-4e23-a567-9737026dd771",
   "metadata": {},
   "source": [
    "### Define Agent Input Component\n",
    "\n",
    "Here we define the agent input component, called at the beginning of every agent step. Besides passing along the input, we also do initialization/state modification."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b1a9252-57df-47cc-9fcf-84d87b1e0102",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.agent.react.types import (\n",
    "    ActionReasoningStep,\n",
    "    ObservationReasoningStep,\n",
    "    ResponseReasoningStep,\n",
    ")\n",
    "from llama_index.agent import Task, AgentChatResponse\n",
    "from llama_index.query_pipeline import (\n",
    "    AgentInputComponent,\n",
    "    AgentFnComponent,\n",
    "    CustomAgentComponent,\n",
    "    ToolRunnerComponent,\n",
    "    QueryComponent,\n",
    ")\n",
    "from llama_index.llms import MessageRole\n",
    "from typing import Dict, Any, Optional, Tuple, List, cast\n",
    "\n",
    "\n",
    "## Agent Input Component\n",
    "## This is the component that produces agent inputs to the rest of the components\n",
    "## Can also put initialization logic here.\n",
    "def agent_input_fn(task: Task, state: Dict[str, Any]) -> Dict[str, Any]:\n",
    "    \"\"\"Agent input function.\n",
    "\n",
    "    Returns:\n",
    "        A Dictionary of output keys and values. If you are specifying\n",
    "        src_key when defining links between this component and other\n",
    "        components, make sure the src_key matches the specified output_key.\n",
    "\n",
    "    \"\"\"\n",
    "    # initialize current_reasoning\n",
    "    if \"current_reasoning\" not in state:\n",
    "        state[\"current_reasoning\"] = []\n",
    "    reasoning_step = ObservationReasoningStep(observation=task.input)\n",
    "    state[\"current_reasoning\"].append(reasoning_step)\n",
    "    return {\"input\": task.input}\n",
    "\n",
    "\n",
    "agent_input_component = AgentInputComponent(fn=agent_input_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd914690-ee4d-4b3b-bb95-9cce0994b9c1",
   "metadata": {},
   "source": [
    "### Define Agent Prompt\n",
    "\n",
    "Here we define the agent component that generates a ReAct prompt, and after the output is generated from the LLM, parses into a structured object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a1957da-bc2b-4186-abf0-11148a23dba7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.agent.react.formatter import ReActChatFormatter\n",
    "from llama_index.query_pipeline import InputComponent, Link\n",
    "from llama_index.llms import ChatMessage\n",
    "from llama_index.tools import BaseTool\n",
    "\n",
    "\n",
    "## define prompt function\n",
    "def react_prompt_fn(\n",
    "    task: Task, state: Dict[str, Any], input: str, tools: List[BaseTool]\n",
    ") -> List[ChatMessage]:\n",
    "    # Add input to reasoning\n",
    "    chat_formatter = ReActChatFormatter()\n",
    "    return chat_formatter.format(\n",
    "        tools,\n",
    "        chat_history=task.memory.get() + state[\"memory\"].get_all(),\n",
    "        current_reasoning=state[\"current_reasoning\"],\n",
    "    )\n",
    "\n",
    "\n",
    "react_prompt_component = AgentFnComponent(\n",
    "    fn=react_prompt_fn, partial_dict={\"tools\": [sql_tool]}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7acd7088-7804-48c5-8c7f-0c5c112d27f0",
   "metadata": {},
   "source": [
    "### Define Agent Output Parser + Tool Pipeline\n",
    "\n",
    "Once the LLM gives an output, we have a decision tree:\n",
    "1. If an answer is given, then we're done. Process the output\n",
    "2. If an action is given, we need to execute the specified tool with the specified args, and then process the output.\n",
    "\n",
    "Tool calling can be done via the `ToolRunnerComponent` module. This is a standalone module that takes in a list of tools, and can be \"executed\" with the specified tool name (every tool has a name) and tool action.\n",
    "\n",
    "We implement this overall module `OutputAgentComponent` that subclasses `CustomAgentComponent`.\n",
    "\n",
    "Note: we also implement `sub_query_components` to pass through higher-level callback managers to the tool runner submodule."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "702e3070-40df-491b-b468-d19ba98edc6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Set, Optional\n",
    "from llama_index.agent.react.output_parser import ReActOutputParser\n",
    "\n",
    "\n",
    "## Agent Output Component\n",
    "## Process reasoning step/tool outputs, and return agent response\n",
    "def finalize_fn(\n",
    "    task: Task,\n",
    "    state: Dict[str, Any],\n",
    "    reasoning_step: Any,\n",
    "    is_done: bool = False,\n",
    "    tool_output: Optional[Any] = None,\n",
    ") -> Tuple[AgentChatResponse, bool]:\n",
    "    \"\"\"Finalize function.\n",
    "\n",
    "    Here we take the latest reasoning step, and a tool output (if provided),\n",
    "    and return the agent output (and decide if agent is done).\n",
    "\n",
    "    This function returns an `AgentChatResponse` and `is_done` tuple. and\n",
    "    is the last component of the query pipeline. This is the expected\n",
    "    return type for any query pipeline passed to `QueryPipelineAgentWorker`.\n",
    "\n",
    "    \"\"\"\n",
    "    current_reasoning = state[\"current_reasoning\"]\n",
    "    current_reasoning.append(reasoning_step)\n",
    "    # if tool_output is not None, add to current reasoning\n",
    "    if tool_output is not None:\n",
    "        observation_step = ObservationReasoningStep(\n",
    "            observation=str(tool_output)\n",
    "        )\n",
    "        current_reasoning.append(observation_step)\n",
    "    if isinstance(current_reasoning[-1], ResponseReasoningStep):\n",
    "        response_step = cast(ResponseReasoningStep, current_reasoning[-1])\n",
    "        response_str = response_step.response\n",
    "    else:\n",
    "        response_str = current_reasoning[-1].get_content()\n",
    "\n",
    "    # if is_done, add to memory\n",
    "    # NOTE: memory is a reserved keyword in `state`, but you can add your own too\n",
    "    if is_done:\n",
    "        state[\"memory\"].put(\n",
    "            ChatMessage(content=task.input, role=MessageRole.USER)\n",
    "        )\n",
    "        state[\"memory\"].put(\n",
    "            ChatMessage(content=response_str, role=MessageRole.ASSISTANT)\n",
    "        )\n",
    "\n",
    "    return AgentChatResponse(response=response_str), is_done\n",
    "\n",
    "\n",
    "class OutputAgentComponent(CustomAgentComponent):\n",
    "    \"\"\"Output agent component.\"\"\"\n",
    "\n",
    "    tool_runner_component: ToolRunnerComponent\n",
    "    output_parser: ReActOutputParser\n",
    "\n",
    "    def __init__(self, tools, **kwargs):\n",
    "        tool_runner_component = ToolRunnerComponent(tools)\n",
    "        super().__init__(\n",
    "            tool_runner_component=tool_runner_component,\n",
    "            output_parser=ReActOutputParser(),\n",
    "            **kwargs\n",
    "        )\n",
    "\n",
    "    def _run_component(self, **kwargs: Any) -> Any:\n",
    "        \"\"\"Run component.\"\"\"\n",
    "        chat_response = kwargs[\"chat_response\"]\n",
    "        task = kwargs[\"task\"]\n",
    "        state = kwargs[\"state\"]\n",
    "        reasoning_step = self.output_parser.parse(\n",
    "            chat_response.message.content\n",
    "        )\n",
    "        if reasoning_step.is_done:\n",
    "            return {\n",
    "                \"output\": finalize_fn(\n",
    "                    task, state, reasoning_step, is_done=True\n",
    "                )\n",
    "            }\n",
    "        else:\n",
    "            tool_output = self.tool_runner_component.run_component(\n",
    "                tool_name=reasoning_step.action,\n",
    "                tool_input=reasoning_step.action_input,\n",
    "            )\n",
    "            return {\n",
    "                \"output\": finalize_fn(\n",
    "                    task,\n",
    "                    state,\n",
    "                    reasoning_step,\n",
    "                    is_done=False,\n",
    "                    tool_output=tool_output,\n",
    "                )\n",
    "            }\n",
    "\n",
    "    @property\n",
    "    def _input_keys(self) -> Set[str]:\n",
    "        return {\"chat_response\"}\n",
    "\n",
    "    @property\n",
    "    def _optional_input_keys(self) -> Set[str]:\n",
    "        return {\"is_done\", \"tool_output\"}\n",
    "\n",
    "    @property\n",
    "    def _output_keys(self) -> Set[str]:\n",
    "        return {\"output\"}\n",
    "\n",
    "    @property\n",
    "    def sub_query_components(self) -> List[QueryComponent]:\n",
    "        return [self.tool_runner_component]\n",
    "\n",
    "\n",
    "react_output_component = OutputAgentComponent([sql_tool])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03c1382b-f9ef-42a6-b117-91d7b2a14921",
   "metadata": {},
   "source": [
    "### Stitch together Agent Query Pipeline\n",
    "\n",
    "We can now stitch together the top-level agent pipeline: agent_input -> react_prompt -> llm -> react_output.\n",
    "\n",
    "The last component is the if-else component that calls sub-components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "219ec2ed-474e-450e-866c-15772055b8b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.query_pipeline import QueryPipeline as QP\n",
    "from llama_index.llms import OpenAI\n",
    "\n",
    "qp = QP(\n",
    "    modules={\n",
    "        \"agent_input\": agent_input_component,\n",
    "        \"react_prompt\": react_prompt_component,\n",
    "        \"llm\": OpenAI(model=\"gpt-4-1106-preview\"),\n",
    "        \"react_output\": react_output_component,\n",
    "    },\n",
    "    verbose=True,\n",
    ")\n",
    "qp.add_chain([\"agent_input\", \"react_prompt\", \"llm\", \"react_output\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "207ff8b0-efc3-482d-8224-f3b82c9a1c2c",
   "metadata": {},
   "source": [
    "### Visualize Query Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd5e739b-ef99-44dd-91a3-a495fa11e4f7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "agent_dag.html\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"600px\"\n",
       "            src=\"agent_dag.html\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x2ac3293f0>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pyvis.network import Network\n",
    "\n",
    "net = Network(notebook=True, cdn_resources=\"in_line\", directed=True)\n",
    "net.from_nx(qp.dag)\n",
    "net.show(\"agent_dag.html\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7785354-1fe0-49b4-8b48-08ce51853a4e",
   "metadata": {},
   "source": [
    "### Setup Agent Worker around Text-to-SQL Query Pipeline\n",
    "\n",
    "This is our way to setup an agent around a text-to-SQL Query Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73d6d178-97e9-4525-82ba-e14030506691",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.agent import QueryPipelineAgentWorker, AgentRunner\n",
    "from llama_index.callbacks import CallbackManager\n",
    "\n",
    "agent_worker = QueryPipelineAgentWorker(qp)\n",
    "agent = AgentRunner(agent_worker, callback_manager=CallbackManager([]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f927ac2e-ef69-4b15-8a10-f1e57ba4aa03",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[AgentFnComponent(partial_dict={'tools': [<llama_index.tools.query_engine.QueryEngineTool object at 0x2ac2c8880>]}, fn=<function react_prompt_fn at 0x2acc2caf0>, async_fn=None),\n",
       " OutputAgentComponent(partial_dict={}, callback_manager=<llama_index.callbacks.base.CallbackManager object at 0x2ac37b5b0>, tool_runner_component=ToolRunnerComponent(partial_dict={}, tool_dict={'sql_tool': <llama_index.tools.query_engine.QueryEngineTool object at 0x2ac2c8880>}, callback_manager=<llama_index.callbacks.base.CallbackManager object at 0x2ac37b070>), output_parser=<llama_index.agent.react.output_parser.ReActOutputParser object at 0x2ac37af20>)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_worker.agent_components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d01ea94b-3d3f-4b26-b71a-eea3d1d1ec0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.agent import QueryPipelineAgentWorker, AgentRunner\n",
    "from llama_index.callbacks import CallbackManager\n",
    "\n",
    "agent_worker = QueryPipelineAgentWorker(qp)\n",
    "agent = AgentRunner(agent_worker, callback_manager=CallbackManager([]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdd8835d-1044-4cca-a6aa-59f7635096ca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[AgentFnComponent(partial_dict={'tools': [<llama_index.tools.query_engine.QueryEngineTool object at 0x2a3e849a0>], 'task': Task(task_id='fe91205c-62ca-4a6f-96dd-c72685b6fa27', input='What are some tracks from the artist AC/DC? Limit it to 3', memory=ChatMemoryBuffer(token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all'), chat_store=SimpleChatStore(store={}), chat_store_key='chat_history'), extra_state={}), 'state': {'sources': [], 'memory': ChatMemoryBuffer(token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all'), chat_store=SimpleChatStore(store={}), chat_store_key='chat_history'), 'current_reasoning': [ObservationReasoningStep(observation='What are some tracks from the artist AC/DC? Limit it to 3'), ActionReasoningStep(thought='I need to use a tool to help me answer the question.', action='sql_tool', action_input={'input': \"SELECT track_name FROM tracks WHERE artist_name = 'AC/DC' LIMIT 3\"}), ObservationReasoningStep(observation='{\\'output\\': ToolOutput(content=\\'Some of AC/DC\\\\\\'s popular tracks are \"For Those About To Rock (We Salute You)\", \"Put The Finger On You\", and \"Let\\\\\\'s Get It Up\".\\', tool_name=\\'sql_tool\\', raw_input={\\'input\\': \"SELECT track_name FROM tracks WHERE artist_name = \\'AC/DC\\' LIMIT 3\"}, raw_output=Response(response=\\'Some of AC/DC\\\\\\'s popular tracks are \"For Those About To Rock (We Salute You)\", \"Put The Finger On You\", and \"Let\\\\\\'s Get It Up\".\\', source_nodes=[NodeWithScore(node=TextNode(id_=\\'0b519588-2fc1-49d9-ac43-d744654ba1dc\\', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, text=\\'[(\\\\\\'For Those About To Rock (We Salute You)\\\\\\',), (\\\\\\'Put The Finger On You\\\\\\',), (\"Let\\\\\\'s Get It Up\",)]\\', start_char_idx=None, end_char_idx=None, text_template=\\'{metadata_str}\\\\n\\\\n{content}\\', metadata_template=\\'{key}: {value}\\', metadata_seperator=\\'\\\\n\\'), score=None)], metadata={\\'0b519588-2fc1-49d9-ac43-d744654ba1dc\\': {}, \\'sql_query\\': \"SELECT tracks.Name AS track_name\\\\nFROM tracks\\\\nJOIN albums ON tracks.AlbumId = albums.AlbumId\\\\nJOIN artists ON albums.ArtistId = artists.ArtistId\\\\nWHERE artists.Name = \\'AC/DC\\'\\\\nLIMIT 3\", \\'result\\': [(\\'For Those About To Rock (We Salute You)\\',), (\\'Put The Finger On You\\',), (\"Let\\'s Get It Up\",)], \\'col_keys\\': [\\'track_name\\']}))}')]}}, fn=<function react_prompt_fn at 0x2a482c820>, async_fn=None),\n",
       " OutputAgentComponent(partial_dict={}, callback_manager=<llama_index.callbacks.base.CallbackManager object at 0x2a3fe7b50>, tool_runner_component=ToolRunnerComponent(partial_dict={}, tool_dict={'sql_tool': <llama_index.tools.query_engine.QueryEngineTool object at 0x2a3e849a0>}, callback_manager=<llama_index.callbacks.base.CallbackManager object at 0x2a3fe7070>), output_parser=<llama_index.agent.react.output_parser.ReActOutputParser object at 0x2a3fe7c70>)]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_worker.agent_components"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "373bd661-4a6b-4c5b-b35b-0082921cc3b9",
   "metadata": {},
   "source": [
    "### Run the Agent\n",
    "\n",
    "Let's try the agent on some sample queries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ea55110-98fb-4acf-b2f8-eee2b15c087b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# start task\n",
    "task = agent.create_task(\n",
    "    \"What are some tracks from the artist AC/DC? Limit it to 3\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f35e816-e76b-43ba-8eb4-ddda8ee97a0b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module agent_input with input: \n",
      "state: {'sources': [], 'memory': ChatMemoryBuffer(token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all'), chat_store=SimpleChatSto...\n",
      "task: task_id='79a8d443-5707-4632-82b2-51fd253cd294' input='What are some tracks from the artist AC/DC? Limit it to 3' memory=ChatMemoryBuffer(token_limit=3000, tokenizer_fn=functools.partial(<bound method ...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module react_prompt with input: \n",
      "input: What are some tracks from the artist AC/DC? Limit it to 3\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module llm with input: \n",
      "messages: [ChatMessage(role=<MessageRole.SYSTEM: 'system'>, content='\\nYou are designed to help with a variety of tasks, from answering questions     to providing summaries to other types of analyses.\\n\\n## Too...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module react_output with input: \n",
      "chat_response: assistant: Thought: I need to use a tool to help me answer the question.\n",
      "Action: sql_tool\n",
      "Action Input: {\"input\": \"Select track_name from music_database where artist_name = 'AC/DC' limit 3\"}\n",
      "\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "step_output = agent.run_step(task.task_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a13e449-db6c-43a3-99b7-70d23687e8b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module agent_input with input: \n",
      "state: {'sources': [], 'memory': ChatMemoryBuffer(token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all'), chat_store=SimpleChatSto...\n",
      "task: task_id='79a8d443-5707-4632-82b2-51fd253cd294' input='What are some tracks from the artist AC/DC? Limit it to 3' memory=ChatMemoryBuffer(token_limit=3000, tokenizer_fn=functools.partial(<bound method ...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module react_prompt with input: \n",
      "input: What are some tracks from the artist AC/DC? Limit it to 3\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module llm with input: \n",
      "messages: [ChatMessage(role=<MessageRole.SYSTEM: 'system'>, content='\\nYou are designed to help with a variety of tasks, from answering questions     to providing summaries to other types of analyses.\\n\\n## Too...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module react_output with input: \n",
      "chat_response: assistant: Thought: The user has repeated the question, but I have already provided the answer using the tool. I will restate the answer.\n",
      "\n",
      "Answer: The top 3 tracks by AC/DC are \"For Those About To Roc...\n",
      "\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "step_output = agent.run_step(task.task_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22dffeec-651e-43e7-8929-827566c8f6da",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "step_output.is_last"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44372a37-0b1e-4d90-8f42-9816546cc66e",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = agent.finalize_response(task.task_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ba842bd-cc86-422b-8c17-2ff87c1962b6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The top 3 tracks by AC/DC are \"For Those About To Rock (We Salute You)\", \"Put The Finger On You\", and \"Let's Get It Up\".\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8f8458f-4fab-420f-bb80-214e1f11c0a8",
   "metadata": {},
   "source": [
    "## Setup Simple Retry Agent Pipeline for Text-to-SQL \n",
    "\n",
    "Instead of the full ReAct pipeline that does tool picking, let's try a much simpler agent pipeline that only does text-to-SQL, with retry-logic.\n",
    "\n",
    "We try a simple text-based \"retry\" prompt where given the user input and previous conversation history, can generate a modified query that outputs the right result."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7adc8aa9-5b72-4525-8bad-052c5df28d6a",
   "metadata": {},
   "source": [
    "### Define Core Modules\n",
    "\n",
    "- agent input\n",
    "- retry prompt\n",
    "- output processor (including a validation prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b62f3c5-6c4d-4532-ab34-fafe9abb77a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms import OpenAI\n",
    "\n",
    "# llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "llm = OpenAI(model=\"gpt-4-1106-preview\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06dd3c71-66d3-4a7d-ba03-abecd61268d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.agent import Task, AgentChatResponse\n",
    "from typing import Dict, Any\n",
    "from llama_index.query_pipeline import AgentInputComponent, AgentFnComponent\n",
    "\n",
    "\n",
    "def agent_input_fn(task: Task, state: Dict[str, Any]) -> Dict:\n",
    "    \"\"\"Agent input function.\"\"\"\n",
    "    # initialize current_reasoning\n",
    "    if \"convo_history\" not in state:\n",
    "        state[\"convo_history\"] = []\n",
    "        state[\"count\"] = 0\n",
    "    state[\"convo_history\"].append(f\"User: {task.input}\")\n",
    "    convo_history_str = \"\\n\".join(state[\"convo_history\"]) or \"None\"\n",
    "    return {\"input\": task.input, \"convo_history\": convo_history_str}\n",
    "\n",
    "\n",
    "agent_input_component = AgentInputComponent(fn=agent_input_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f18356fa-df3b-4fbf-ab5b-c2708b0bfaa4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.prompts import PromptTemplate\n",
    "\n",
    "retry_prompt_str = \"\"\"\\\n",
    "You are trying to generate a proper natural language query given a user input.\n",
    "\n",
    "This query will then be interpreted by a downstream text-to-SQL agent which\n",
    "will convert the query to a SQL statement. If the agent triggers an error,\n",
    "then that will be reflected in the current conversation history (see below).\n",
    "\n",
    "If the conversation history is None, use the user input. If its not None,\n",
    "generate a new SQL query that avoids the problems of the previous SQL query.\n",
    "\n",
    "Input: {input}\n",
    "Convo history (failed attempts): \n",
    "{convo_history}\n",
    "\n",
    "New input: \"\"\"\n",
    "retry_prompt = PromptTemplate(retry_prompt_str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f14e034c-b5b3-4f7d-b840-b7f2f4926d79",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.response import Response\n",
    "from typing import Tuple\n",
    "\n",
    "validate_prompt_str = \"\"\"\\\n",
    "Given the user query, validate whether the inferred SQL query and response from executing the query is correct and answers the query.\n",
    "\n",
    "Answer with YES or NO.\n",
    "\n",
    "Query: {input}\n",
    "Inferred SQL query: {sql_query}\n",
    "SQL Response: {sql_response}\n",
    "\n",
    "Result: \"\"\"\n",
    "validate_prompt = PromptTemplate(validate_prompt_str)\n",
    "\n",
    "MAX_ITER = 3\n",
    "\n",
    "\n",
    "def agent_output_fn(\n",
    "    task: Task, state: Dict[str, Any], output: Response\n",
    ") -> Tuple[AgentChatResponse, bool]:\n",
    "    \"\"\"Agent output component.\"\"\"\n",
    "    print(f\"> Inferred SQL Query: {output.metadata['sql_query']}\")\n",
    "    print(f\"> SQL Response: {str(output)}\")\n",
    "    state[\"convo_history\"].append(\n",
    "        f\"Assistant (inferred SQL query): {output.metadata['sql_query']}\"\n",
    "    )\n",
    "    state[\"convo_history\"].append(f\"Assistant (response): {str(output)}\")\n",
    "\n",
    "    # run a mini chain to get response\n",
    "    validate_prompt_partial = validate_prompt.as_query_component(\n",
    "        partial={\n",
    "            \"sql_query\": output.metadata[\"sql_query\"],\n",
    "            \"sql_response\": str(output),\n",
    "        }\n",
    "    )\n",
    "    qp = QP(chain=[validate_prompt_partial, llm])\n",
    "    validate_output = qp.run(input=task.input)\n",
    "\n",
    "    state[\"count\"] += 1\n",
    "    is_done = False\n",
    "    if state[\"count\"] >= MAX_ITER:\n",
    "        is_done = True\n",
    "    if \"YES\" in validate_output.message.content:\n",
    "        is_done = True\n",
    "\n",
    "    return AgentChatResponse(response=str(output)), is_done\n",
    "\n",
    "\n",
    "agent_output_component = AgentFnComponent(fn=agent_output_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44015fad-a37d-48e7-90d0-77d62ff4ab83",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.query_pipeline import (\n",
    "    QueryPipeline as QP,\n",
    "    Link,\n",
    "    InputComponent,\n",
    ")\n",
    "\n",
    "qp = QP(\n",
    "    modules={\n",
    "        \"input\": agent_input_component,\n",
    "        \"retry_prompt\": retry_prompt,\n",
    "        \"llm\": llm,\n",
    "        \"sql_query_engine\": sql_query_engine,\n",
    "        \"output_component\": agent_output_component,\n",
    "    },\n",
    "    verbose=True,\n",
    ")\n",
    "qp.add_link(\"input\", \"retry_prompt\", src_key=\"input\", dest_key=\"input\")\n",
    "qp.add_link(\n",
    "    \"input\", \"retry_prompt\", src_key=\"convo_history\", dest_key=\"convo_history\"\n",
    ")\n",
    "qp.add_chain([\"retry_prompt\", \"llm\", \"sql_query_engine\", \"output_component\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64185c9e-b90f-4acb-b498-64a1bb4f8751",
   "metadata": {},
   "source": [
    "### Visualize Query Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "142dd638-578c-42c0-aa06-395052ca210a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "agent_dag.html\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"600px\"\n",
       "            src=\"agent_dag.html\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x2ac2c8a30>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pyvis.network import Network\n",
    "\n",
    "net = Network(notebook=True, cdn_resources=\"in_line\", directed=True)\n",
    "net.from_nx(qp.dag)\n",
    "net.show(\"agent_dag.html\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab09a855-ef74-457c-8aa6-2292adf12278",
   "metadata": {},
   "source": [
    "### Define Agent Worker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3764e8c-78a3-41cd-97f4-8c6c1642ff1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.agent import QueryPipelineAgentWorker, AgentRunner\n",
    "from llama_index.callbacks import CallbackManager\n",
    "\n",
    "agent_worker = QueryPipelineAgentWorker(qp)\n",
    "agent = AgentRunner(agent_worker, callback_manager=CallbackManager([]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d9a327a-e9ee-481e-b427-701245cdd2f1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module input with input: \n",
      "state: {'sources': [], 'memory': ChatMemoryBuffer(token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all'), chat_store=SimpleChatSto...\n",
      "task: task_id='741c0d59-fa40-44a2-acab-cc4c36fdf0c7' input=\"How many albums did the artist who wrote 'Restless and Wild' release? (answer should be non-zero)?\" memory=ChatMemoryBuffer(token_limit=3000, toke...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module retry_prompt with input: \n",
      "input: How many albums did the artist who wrote 'Restless and Wild' release? (answer should be non-zero)?\n",
      "convo_history: User: How many albums did the artist who wrote 'Restless and Wild' release? (answer should be non-zero)?\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module llm with input: \n",
      "messages: You are trying to generate a proper natural language query given a user input.\n",
      "\n",
      "This query will then be interpreted by a downstream text-to-SQL agent which\n",
      "will convert the query to a SQL statement. I...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module sql_query_engine with input: \n",
      "input: assistant: Given the user input and the requirement that the answer should be non-zero, the natural language query should be specific enough to avoid ambiguity and errors when converted to SQL. The pr...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module output_component with input: \n",
      "output: The count of albums released by the artist credited with writing the song titled 'Restless and Wild', excluding any counts of zero, is 1.\n",
      "\n",
      "\u001b[0m> Inferred SQL Query: SELECT COUNT(*) \n",
      "FROM albums \n",
      "JOIN artists ON albums.ArtistId = artists.ArtistId \n",
      "JOIN tracks ON albums.AlbumId = tracks.AlbumId \n",
      "WHERE tracks.Name = 'Restless and Wild' \n",
      "AND albums.ArtistId IS NOT NULL \n",
      "AND albums.AlbumId IS NOT NULL \n",
      "AND albums.Title IS NOT NULL \n",
      "AND artists.Name IS NOT NULL \n",
      "AND tracks.TrackId IS NOT NULL \n",
      "AND tracks.MediaTypeId IS NOT NULL \n",
      "AND tracks.GenreId IS NOT NULL \n",
      "AND tracks.Composer IS NOT NULL \n",
      "AND tracks.Milliseconds IS NOT NULL \n",
      "AND tracks.Bytes IS NOT NULL \n",
      "AND tracks.UnitPrice IS NOT NULL\n",
      "> SQL Response: The count of albums released by the artist credited with writing the song titled 'Restless and Wild', excluding any counts of zero, is 1.\n",
      "\u001b[1;3;38;2;155;135;227m> Running module input with input: \n",
      "state: {'sources': [], 'memory': ChatMemoryBuffer(token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all'), chat_store=SimpleChatSto...\n",
      "task: task_id='741c0d59-fa40-44a2-acab-cc4c36fdf0c7' input=\"How many albums did the artist who wrote 'Restless and Wild' release? (answer should be non-zero)?\" memory=ChatMemoryBuffer(token_limit=3000, toke...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module retry_prompt with input: \n",
      "input: How many albums did the artist who wrote 'Restless and Wild' release? (answer should be non-zero)?\n",
      "convo_history: User: How many albums did the artist who wrote 'Restless and Wild' release? (answer should be non-zero)?\n",
      "Assistant (inferred SQL query): SELECT COUNT(*) \n",
      "FROM albums \n",
      "JOIN artists ON albums.ArtistId =...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module llm with input: \n",
      "messages: You are trying to generate a proper natural language query given a user input.\n",
      "\n",
      "This query will then be interpreted by a downstream text-to-SQL agent which\n",
      "will convert the query to a SQL statement. I...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module sql_query_engine with input: \n",
      "input: assistant: Given the previous failed attempt, it seems that the SQL query was overly complex and included unnecessary conditions. The query should focus on finding the artist who wrote 'Restless and W...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module output_component with input: \n",
      "output: The number of albums released by the artist who composed the track 'Restless and Wild' is 347.\n",
      "\n",
      "\u001b[0m> Inferred SQL Query: SELECT COUNT(*) \n",
      "FROM albums \n",
      "WHERE ArtistId = (SELECT ArtistId \n",
      "                  FROM tracks \n",
      "                  WHERE Name = 'Restless and Wild')\n",
      "> SQL Response: The number of albums released by the artist who composed the track 'Restless and Wild' is 347.\n",
      "\u001b[1;3;38;2;155;135;227m> Running module input with input: \n",
      "state: {'sources': [], 'memory': ChatMemoryBuffer(token_limit=3000, tokenizer_fn=functools.partial(<bound method Encoding.encode of <Encoding 'cl100k_base'>>, allowed_special='all'), chat_store=SimpleChatSto...\n",
      "task: task_id='741c0d59-fa40-44a2-acab-cc4c36fdf0c7' input=\"How many albums did the artist who wrote 'Restless and Wild' release? (answer should be non-zero)?\" memory=ChatMemoryBuffer(token_limit=3000, toke...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module retry_prompt with input: \n",
      "input: How many albums did the artist who wrote 'Restless and Wild' release? (answer should be non-zero)?\n",
      "convo_history: User: How many albums did the artist who wrote 'Restless and Wild' release? (answer should be non-zero)?\n",
      "Assistant (inferred SQL query): SELECT COUNT(*) \n",
      "FROM albums \n",
      "JOIN artists ON albums.ArtistId =...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module llm with input: \n",
      "messages: You are trying to generate a proper natural language query given a user input.\n",
      "\n",
      "This query will then be interpreted by a downstream text-to-SQL agent which\n",
      "will convert the query to a SQL statement. I...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module sql_query_engine with input: \n",
      "input: assistant: Given the previous failed attempts and the user's insistence on a non-zero answer, it seems that the SQL queries might have been correct in structure but the responses provided were not sat...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module output_component with input: \n",
      "output: The number of albums released by the composer of the track 'Restless and Wild' is 1.\n",
      "\n",
      "\u001b[0m> Inferred SQL Query: SELECT COUNT(DISTINCT albums.AlbumId) AS AlbumCount\n",
      "FROM albums\n",
      "JOIN tracks ON albums.AlbumId = tracks.AlbumId\n",
      "WHERE tracks.Composer IN (SELECT Composer FROM tracks WHERE Name = 'Restless and Wild')\n",
      "AND albums.AlbumId IS NOT NULL\n",
      "HAVING AlbumCount > 0;\n",
      "> SQL Response: The number of albums released by the composer of the track 'Restless and Wild' is 1.\n",
      "The number of albums released by the composer of the track 'Restless and Wild' is 1.\n"
     ]
    }
   ],
   "source": [
    "response = agent.chat(\n",
    "    \"How many albums did the artist who wrote 'Restless and Wild' release? (answer should be non-zero)?\"\n",
    ")\n",
    "print(str(response))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_index_v2",
   "language": "python",
   "name": "llama_index_v2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
